import hashlib
import random


class Merkle_Tree:
    def list_hashing(self, data_list):
        """
        对数据列表中的每个数据进行哈希处理
        :param data_list: 数据列表
        :return: 哈希值列表
        """
        result_hash_list = list()
        for data in data_list:
            sha_256 = hashlib.sha256()
            sha_256.update(str(data).encode('utf-8'))
            result_hash_list.append(sha_256.hexdigest())
        return result_hash_list

    def catenate_hash(self, data1, data2):
        """
        将两个哈希值拼接后进行哈希处理
        :param data1: 第一个哈希值
        :param data2: 第二个哈希值
        :return: 拼接后的哈希值
        """
        sha_256 = hashlib.sha256()
        sha_256.update(data1.encode("utf-8"))
        sha_256.update(data2.encode("utf-8"))
        return sha_256.hexdigest()

    def __init__(self, data_list):
        """
        初始化Merkle树
        :param data_list: 数据列表
        """
        self.data_list = data_list
        self._merkle_tree = dict()
        self.hash_list = self.list_hashing(data_list)
        mt, self.layer_length = self.generate(self.hash_list)
        self.hash_list = mt
        last_layer_index = len(self.layer_length) - 1
        self.merkle_root = self.hash_list[f"layer_{last_layer_index}"]

    def generate(self, data_list):
        """
        生成Merkle树
        :param data_list: 哈希值列表
        :return: Merkle树的哈希值列表和每层的节点数
        """
        isOdd = True
        if len(data_list) % 2 == 0:
            isOdd = False
        pair_num = int(len(data_list) / 2)
        mt = dict()
        layer_count = 0
        layer_list = list()
        layer_length = list()
        while len(layer_list) >= 1 or layer_count == 0:
            if layer_count == 0:
                layer_list = data_list
            mt[f"layer_{layer_count}"] = layer_list
            layer_length.append(len(layer_list))
            layer_count += 1
            if len(layer_list) == 1:
                break
            next_layer_list = list()
            for i in range(pair_num):
                next_layer_list.append(self.catenate_hash(layer_list[2 * i], layer_list[2 * i + 1]))
            if isOdd:
                next_layer_list.append(self.catenate_hash(layer_list[-1], layer_list[-1]))
            layer_list = next_layer_list
            isOdd = True
            if len(layer_list) % 2 == 0:
                isOdd = False
            pair_num = int(len(layer_list) / 2)
        return mt, layer_length


# 测试代码
print("[Merkle树测试]")
data_count = random.randint(8, 20)
data = list()
for i in range(data_count):
    data.append(f"交易_{random.randint(0, 200)}")
print(f"本区块中原始共{len(data)}个交易数据：")
print(data)
mt = Merkle_Tree(data)
print("Merkle树：", mt.hash_list)
layer_count = len(mt.layer_length)
print(f"Merkle Tree层数为：{layer_count}")
print("Merkle Tree各层节点数为：", mt.layer_length)
print(f"Merkle Tree的Top Hash节点值：{mt.merkle_root}")